import os
from datetime import datetime
import xarray as xr
import numpy as np

from config import conf
from create_nc import generate_new_nc_file
from log_own import logger


def produce_2t(file_list):
    global lon_array, lat_array, time_units, init_time
    t2m_list = []
    t_list = []

    for file in file_list:
        try:
            if file.name.endswith("enfo-ef.grib2") or file.name.endswith("enfo-ep.grib2"):
                return
            logger.info("开始生产 " + file.name + " 的数据")
            ds = xr.open_dataset(file, filter_by_keys={'shortName': "2t"})
            da = ds["t2m"]

            # 取当前高度层的数据
            data = da.data
            # 获取 起报时间 和 预报时间
            init_time = file.stem.split('-')[0]  # 20230612000000
            time_units = datetime.strptime(init_time, "%Y%m%d%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
            forecast_time = np.int16(file.stem.split('-')[1].replace("h", ""))  # 0
            # 获取经纬度数组
            lon_array = da["longitude"].data
            lat_array = da["latitude"].data

            t2m_list.append(data)
            t_list.append(forecast_time)
        except EOFError:
            logger.error(str(file) + "生产失败")

    if len(t2m_list) == 0:
        return

    t2m_array = np.array(t2m_list)
    time_array = np.array(t_list)

    variables = {
        'longitude': {
            'array': lon_array,
            'nodata': -99999,
        },
        'latitude': {
            'array': lat_array,
            'nodata': -99999
        },
        'time': {
            'array': time_array,
            'units': "hours since " + time_units
        },
        '2t': {
            'array': t2m_array,
            'nodata': -99999,
            'real_name': '2t'
        }
    }
    store_dir = os.path.join(conf["store_nc_dir"], "2t", init_time)
    if not os.path.exists(store_dir):
        os.makedirs(store_dir)
    nc_save_path = os.path.join(store_dir, ("2t" + ".nc"))
    if not os.path.exists(nc_save_path):
        # 生成nc文件
        generate_new_nc_file(variables, nc_save_path)